"""


Version: 0.1
Author: lk
Date: 2022-03-08 12:11
"""
import dgl
import dgl.function as fn
import torch
import torch.nn as nn
import torch.nn.functional as F
# from dgl.model_zoo.chem.gnn import GCNLayer
from dgl.nn.pytorch import GraphConv
GCNLayer=GraphConv
class UnsupervisedGCN(nn.Module):
    def __init__(self, hidden_size=32, num_layer=2):
        super(UnsupervisedGCN, self).__init__()
        self.layers = nn.ModuleList(
            [
                GCNLayer(
                    hidden_size, hidden_size,activation= F.relu if i + 1 < num_layer else nn.Sequential()
                )
                for i in range(num_layer)
            ]
        )
    def forward(self, g: dgl.DGLGraph, feature):
        for layer in self.layers:
            feats = layer(g, feature)
        return feats


if __name__ == '__main__':
    # g: dgl.DGLGraph
    # g = dgl.graph((torch.tensor([0, 1, 2, 3]),
    #                torch.tensor([1, 2, 3, 4])))
    # g.ndata["h"] = torch.ones(5, 2)
    #
    # model = GCN(hidden_size=2)
    # # print(model)
    # # gcn_layer=GCNLayer(in_feats=2,out_feats=8,activation=None)
    # # print(gcn_layer(g,g.ndata["h"]))
    # out = model(g, g.ndata["h"])
    # print(out )
    #
    # # def f(node):
    # #     # return node.data["h"] * 2
    # #     return {"x": node.data["h"] * 2}
    #
    # # fn.copy_u()
    # #
    # # g.update_all(message_func=message_fn, reduce_func=reduce_fn)
    # # g.apply_nodes(f)
    # # print(g.ndata)
    # # print(type(g.ndata))
    # # print(g.ndata.pop("h"))
    # # print(g.ndata)
    model = UnsupervisedGCN()
    print(model)
    g = dgl.DGLGraph()
    g.add_nodes(3)
    g=dgl.add_self_loop(g)
    g.add_edges([0, 0, 1], [1, 2, 2])
    feat = torch.rand(3, 32)
    # print(model(g, feat).shape)
    print(model(g,feat))
